import SimpleITK as sitk
import matplotlib.pyplot as plt
import numpy
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import unary_from_labels, create_pairwise_gaussian
from pydensecrf.utils import unary_from_softmax, create_pairwise_bilateral
from scipy import ndimage
from skimage import filters
from tqdm import tqdm

from utils.DataUtils import preprocess


def closing(input_image):
    struct_s = ndimage.generate_binary_structure(3, 2)
    result = ndimage.binary_closing(input_image, structure=struct_s, iterations=1)
    result = result.astype(numpy.int8)
    return result




def crf(original_image, predict_image):
    crf_result = np.zeros_like(predict_image)
    n_labels = 2
    image_shape = predict_image.shape

    # 使用densecrf类
    d = dcrf.DenseCRF(image_shape[0] * image_shape[1] * image_shape[2], n_labels)

    # 得到一元势（负对数概率）
    U = unary_from_labels(predict_image, n_labels, gt_prob=0.5, zero_unsure=None)
    # U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=HAS_UNK)## 如果有不确定区域，用这一行代码替换上一行
    d.setUnaryEnergy(U)

    # 这将创建与颜色无关的功能，然后将它们添加到CRF中
    feats = create_pairwise_gaussian(sdims=(3, 3, 3), shape=image_shape)
    d.addPairwiseEnergy(feats, compat=8, kernel=dcrf.DIAG_KERNEL,
                        normalization=dcrf.NORMALIZE_SYMMETRIC)

    # 这将创建与颜色相关的功能，然后将它们添加到CRF中
    feats = create_pairwise_bilateral(sdims=(40, 40, 40), schan=13,
                                      img=original_image, chdim=-1)
    d.addPairwiseEnergy(feats, compat=10,
                        kernel=dcrf.DIAG_KERNEL,
                        normalization=dcrf.NORMALIZE_SYMMETRIC)
    # 进行5次推理
    Q = d.inference(10)

    # 找出每个像素最可能的类
    MAP = np.argmax(Q, axis=0)

    # 将predicted_image转换回相应的颜色并保存图像
    # MAP = colorize[MAP, :]
    crf_result = MAP.reshape(shape=image_shape)
    return crf_result


def dense_crf(original_image, predict_image):
    H = predict_image.shape[0]
    W = predict_image.shape[2]
    if_show = False
    class_num = 2
    result_3 = np.zeros_like(predict_image)
    for i in tqdm(range(107, 338)):
        # i = 160
        img_1 = predict_image[:, i, :]
        img_2 = 1 - predict_image[:, i, :]
        probs = np.stack([img_1, img_2], axis=0)

        if if_show:
            plt.figure(figsize=(15, 5))
            plt.subplot(1, 2, 1)
            plt.imshow(img_1)
            plt.title('Foreground probability')
            plt.axis('off')
            plt.colorbar()
            plt.subplot(1, 2, 2)
            plt.imshow(img_2)
            plt.title('Background probability')
            plt.axis('off')
            plt.colorbar()

        U = unary_from_softmax(probs)  # note: num classes is first dim
        d = dcrf.DenseCRF2D(W, H, class_num)
        d.setUnaryEnergy(U)

        # NCHAN = 1
        #
        # img = np.zeros((H, W, NCHAN), np.uint8)
        # img[H // 3:2 * H // 3, W // 4:3 * W // 4, :] = 1
        #
        # if if_show:
        #     plt.imshow(img[:, :, 0])
        #     plt.title('Bilateral image')
        #     plt.axis('off')
        #     plt.colorbar()
        img = original_image[:, i, :]
        pairwise_energy = create_pairwise_bilateral(sdims=(10, 10), schan=(0.01,),
                                                    img=img, chdim=-1)

        # pairwise_energy now contains as many dimensions as the DenseCRF has features,
        # which in this case is 3: (x,y,channel1)
        img_en = pairwise_energy.reshape((-1, H, W))  # Reshape just for plotting
        if if_show:
            plt.figure(figsize=(15, 5))
            plt.subplot(1, 3, 1)
            plt.imshow(img_en[0])
            plt.title('Pairwise bilateral [x]')
            plt.axis('off')
            plt.colorbar()
            plt.subplot(1, 3, 2)
            plt.imshow(img_en[1])
            plt.title('Pairwise bilateral [y]')
            plt.axis('off')
            plt.colorbar()
            plt.subplot(1, 3, 3)
            plt.imshow(img_en[2])
            plt.title('Pairwise bilateral [c]')
            plt.axis('off')
            plt.colorbar()

        # d = dcrf.DenseCRF2D(W, H, class_num)
        # d.setUnaryEnergy(U)
        d.addPairwiseEnergy(pairwise_energy, compat=10)  # `compat` is the "strength" of this potential.

        # This time, let's do inference in steps ourselves
        # so that we can look at intermediate solutions
        # as well as monitor KL-divergence, which indicates
        # how well we have converged.
        # PyDenseCRF also requires us to keep track of two
        # temporary buffers it needs for computations.
        Q, tmp1, tmp2 = d.startInference()
        # for _ in range(5):
        #     d.stepInference(Q, tmp1, tmp2)
        # kl1 = d.klDivergence(Q) / (H * W)
        # map_soln1 = np.argmax(Q, axis=0).reshape((H, W))
        #
        # for _ in range(20):
        #     d.stepInference(Q, tmp1, tmp2)
        # kl2 = d.klDivergence(Q) / (H * W)
        # map_soln2 = np.argmax(Q, axis=0).reshape((H, W))

        for _ in range(50):
            d.stepInference(Q, tmp1, tmp2)
        kl3 = d.klDivergence(Q) / (H * W)
        map_soln3 = np.argmax(Q, axis=0).reshape((H, W))
        # result_1[:, :, i] = map_soln1
        # result_2[:, :, i] = map_soln2
        result_3[:, i, :] = map_soln3

        if if_show:
            img_en = pairwise_energy.reshape((-1, H, W))  # Reshape just for plotting
            # plt.figure(figsize=(15, 5))
            # plt.subplot(1, 3, 1)
            # plt.imshow(map_soln1)
            # plt.title(r'MAP Solution with DenseCRF\n(5 steps, KL={:.2f})'.format(kl1))
            # plt.axis('off')
            # plt.subplot(1, 3, 2)
            # plt.imshow(map_soln2)
            # plt.title(r'MAP Solution with DenseCRF\n(20 steps, KL={:.2f})'.format(kl2))
            # plt.axis('off')
            plt.subplot(1, 3, 3)
            plt.imshow(map_soln3)
            plt.title(r'MAP Solution with DenseCRF\n(75 steps, KL={:.2f})'.format(kl3))
            plt.axis('off')
        print("good")
    return result_3


def max_connected_domain(predict_image, max_region_num=5):
    predict_image_itk = sitk.GetImageFromArray(predict_image)
    my_filter = sitk.ConnectedComponentImageFilter()
    my_filter.SetFullyConnected(True)
    output = my_filter.Execute(predict_image_itk)
    np_output = sitk.GetArrayFromImage(output)

    lss_filter = sitk.LabelShapeStatisticsImageFilter()
    lss_filter.Execute(output)

    connected_region_num = my_filter.GetObjectCount()

    region_index_dict = {}

    for i in range(1, connected_region_num):  # Label 0 is background
        curr_area = lss_filter.GetNumberOfPixels(i)
        region_index_dict[i] = curr_area

    top_k_region = {k: v for k, v in sorted(region_index_dict.items(), key=lambda items: items[1], reverse=True)}
    tok_k_list = [k for k in top_k_region.keys()][0:max_region_num]

    result = numpy.zeros_like(predict_image)
    for k in tok_k_list:
        result[np_output == k] = 1

    return result


if __name__ == '__main__':
    Threshold = 0.5
    # test_img = np.random.rand(64, 64, 64)
    # test_img = np.where(test_img > Threshold, 1, 0)
    # test_result = closing(test_img)
    # print("complete")
    test_img_path = r"D:\dataset\vessel\Parse_2022_train_data\PA000309"
    from utils.DataUtils import read_nii_file_with_mask
    from models.EvaluationMetrics import *

    train_img, label, train_I_affine, label_I_affine, Box_boundary = read_nii_file_with_mask(test_img_path, )
    original = preprocess(train_img)

    # original = np.where(original > Threshold, 1, 0)
    predict = np.load(r"H:\laboratory\project\blood_vessel\MICCAI2022-Draft\main\predict_result_threshold.npy")
    predict = np.where(predict > Threshold, 1, 0)
    dice = dice_coefficient_3D(predict, label, 1).item()
    # HD = Hausdorff_Distance_3D(predict, label)
    print("The whole dice is " + str(dice))
    # print("The whole HD is " + str(HD))

    # post_processing_result = dense_crf(original, predict)
    # np.save("crf_result", post_processing_result)
    post_processing_result = max_connected_domain(predict)
    np.save("max_region_connect_result", post_processing_result)

    # post_processing_result = torch.from_numpy(post_processing_result).unsqueeze(0).unsqueeze(0)
    whole_dice = dice_coefficient_3D(post_processing_result, label, 1).item()
    # whole_HD = Hausdorff_Distance_3D(post_processing_result, label)
    print("After dense_crf operation, the whole dice is " + str(whole_dice))
    # print("After dense_crf operation, the whole HD is " + str(whole_HD))

    print("good")
